package com.zt.bigdata.elasticsearch.utils.transport;

import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.transport.TransportClient;
import org.elasticsearch.cluster.node.DiscoveryNode;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.transport.TransportAddress;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.aggregations.AggregationBuilder;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.bucket.filter.Filters;
import org.elasticsearch.search.aggregations.bucket.filter.FiltersAggregator;
import org.elasticsearch.search.aggregations.metrics.tophits.TopHits;
import org.elasticsearch.search.sort.SortOrder;
import org.elasticsearch.transport.client.PreBuiltTransportClient;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.InetAddress;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;


/**
 * @author zhangtong
 */
public class ESClient {
    private static final Logger logger = LoggerFactory.getLogger(ESClient.class);
    private TransportClient client;

    public ESClient(String clusterName, String clusterNode) {
        try {
            Settings settings = Settings.builder()
                    .put("cluster.name", clusterName)
                    .build();
            client = new PreBuiltTransportClient(settings);
            String[] hostPorts = clusterNode.split(",");
            for (int i = 0; i < hostPorts.length; i++) {
                String[] hostPort = hostPorts[i].split(":");
                client.addTransportAddress(
                        new TransportAddress(InetAddress.getByName(hostPort[0]), Integer.valueOf(hostPort[1])));
            }

            List<DiscoveryNode> discoveryNodes = client.connectedNodes();
            if (discoveryNodes.size() > 0) {
                discoveryNodes.forEach(
                        discoveryNode -> logger.info(discoveryNode.getHostAddress())
                );
            } else {
                logger.error("connect failed");
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public TransportClient getClient() {
        return client;
    }

    public static void main(String[] args) {
        new ESClient("my-application","172.28.15.78:9300");

    }

    public List<Map<String, Object>> matchAll() {
        List<Map<String, Object>> list = new ArrayList<>();
        SearchResponse searchResponse = client.prepareSearch("my-bank")
                .setTypes("_doc")
                .setSize(10000)
                //这个游标维持多长时间
                .setScroll(TimeValue.timeValueMinutes(8))
                .execute()
                .actionGet();

        while (true) {
            for (SearchHit searchHit : searchResponse.getHits()) {
                list.add(searchHit.getSourceAsMap());
            }
            searchResponse = client.prepareSearchScroll(searchResponse.getScrollId())
                    .setScroll(TimeValue.timeValueMinutes(8))
                    .execute()
                    .actionGet();

            if (searchResponse.getHits().getHits().length == 0) {
                break;
            }
        }
        return list;
    }

    public String index(String index, String type, String id, Map<String, Object> source) {
        IndexResponse response
                = client.prepareIndex(index, type, id)
                .setSource(source)
                .get();

        return response.getId();
    }

    public Map<String, Object> getDoc(String index, String type, String id) {
        GetResponse response = client.prepareGet(index, type, id).get();
        return response.getSource();
    }

    public Integer deleteDoc(String index, String type, String id) {
        DeleteResponse response = client.prepareDelete(index, type, id).get();
        return response.status().getStatus();
    }

    /**
     * 男女富豪榜
     */
    public Map<String, List<Map<String, Object>>> fuhaobang(int size) {

        Map<String, List<Map<String, Object>>> map = new HashMap<>();
        long start = System.currentTimeMillis();
        AggregationBuilder aggregation =
                AggregationBuilders
                        .filters("agg",
                                new FiltersAggregator.KeyedFilter("men", QueryBuilders.termQuery("gender", "m")),
                                new FiltersAggregator.KeyedFilter("women", QueryBuilders.termQuery("gender", "f")));

        SearchResponse searchResponse = client.prepareSearch("bank2")
                .setTypes("_doc")
                .addAggregation(aggregation
                        .subAggregation(AggregationBuilders
                                .topHits("top")
                                .explain(true)
                                .sort("balance", SortOrder.DESC)
                                .size(size)))
                .get();
        long end = System.currentTimeMillis();
        System.out.println("time used: " + (end - start));

        Filters agg = searchResponse.getAggregations().get("agg");

        for (Filters.Bucket entry : agg.getBuckets()) {
            List<Map<String, Object>> list = new ArrayList<>();
            String key = entry.getKeyAsString();            // bucket key
            long docCount = entry.getDocCount();            // Doc count
            System.out.println("key [{" + key + "}], doc_count [{" + docCount + "}]");
            TopHits topHits = entry.getAggregations().get("top");
            for (SearchHit hit : topHits.getHits().getHits()) {
                list.add(hit.getSourceAsMap());
                System.out.println(" -> id [{" + hit.getId() + "}], _source [{" + hit.getSourceAsString() + "}]");
            }
            map.put(key, list);
        }
        return map;
    }
}
